import os
import re
import datetime
from . import settings

def writeinputfile(VTLines):
    INCAR=[]
    KPOINTS=[]
    VTValues={}
    KPTAG=False
    for Line in VTLines:
        ROBOOTLine=Line.strip()
        #如果是定义的KPOINTS的开头，则标记TAG：
        if ROBOOTLine[0:6]=='%START':
            KPTAG=True
            continue
        #如果定义的是KPOINTS的结尾，则标记TAG：
        if ROBOOTLine[0:4]=="%END":
            KPTAG=False
            continue        
        #如果是KPOINTS行内的东西，则放到KPOINTS列表里
        if KPTAG:
            KPOINTS.append(Line)
            continue
        #如果是写在INCAR里的注释，则原样放上去
        if ROBOOTLine[0]=='#':
            INCAR.append(Line)
            continue
        UNOTELine=ROBOOTLine.split('#')[0]        
        #如果是定义VT参数的值，则记录该值
        if '=' in UNOTELine and UNOTELine[0]=='%':
            VTName=UNOTELine.split('=')[0].strip().split('%')[1].strip()
            VTValue=UNOTELine.split('=')[1].strip()
            VTValues[VTName]=VTValue
            continue
        #其他东西也不管了，一股脑塞INCAR里
        INCAR.append(Line)
    CALCDIR=WriteKPOINTS(KPOINTS,VTValues)
    POTCARINFO=WritePOTCAR(VTValues,CALCDIR)
    WriteINCAR(INCAR,VTValues,POTCARINFO,CALCDIR)

def WriteKPOINTS(KPOINTS,VTValues):
    #先建个文件夹
    if "CALCDIR" in VTValues:
        DIR=VTValues["CALCDIR"]
    else:
        DIR="CALC"
    os.makedirs(DIR)
    fp=open(os.path.join(DIR,"KPOINTS"),'w')
    fp.writelines(KPOINTS)
    fp.close()
    return DIR

def WritePOTCAR(VTValues,CALCDIR):
    #先判断赝势文件目录规定了没
    if not "PSEUDO_DIR" in VTValues:
        raise Exception("You need %PSEUDO_DIR to tell us the root dir of your pseudopotential files")
    else:
        PSEUDO_DIR=VTValues["PSEUDO_DIR"]
    #再瞅瞅有没有指定POSCAR文件
    if "POSCAR" in VTValues:
        PSFile=VTValues["POSCAR"]
    else:
        PSFile="POSCAR"
    #判断POSCAR文件是否可以访问
    if not os.access(PSFile,os.R_OK):
        raise Exception("POSCAR file can not read")
    POSfp=open(PSFile,"r")
    POSCAR=POSfp.readlines()
    POSfp.close()
    POSfp=open(os.path.join(CALCDIR,"POSCAR"),'w')
    POSfp.writelines(POSCAR)
    POSfp.close()
    open(os.path.join(CALCDIR,"POTCAR"),'w').close()
    Atoms=POSCAR[5].split()
    POTCARS=[]
    for Atom in Atoms:
        if "POT-"+Atom in VTValues:
            PTFile=VTValues["POT-"+Atom]
        else:
            PTFile=Atom
        if not os.access(os.path.join(PSEUDO_DIR,PTFile),os.R_OK):
            raise Exception("POSCAR file can not read:"+os.path.join(PSEUDO_DIR,PTFile))
        else:
            AtomDir=os.path.join(PSEUDO_DIR,PTFile)
        if os.access(os.path.join(AtomDir,"POTCAR"),os.R_OK):
            os.system("cat %s >> %s"%(os.path.join(AtomDir,"POTCAR"),os.path.join(CALCDIR,"POTCAR")))
            POTCARS.append((Atom,os.path.join(AtomDir,"POTCAR")))
        elif os.access(os.path.join(AtomDir,"POTCAR.Z"),os.R_OK):
            os.system("zcat %s >> %s"%(os.path.join(AtomDir,"POTCAR.Z"),os.path.join(CALCDIR,"POTCAR")))
            POTCARS.append((Atom,os.path.join(AtomDir,"POTCAR.Z")))
        elif os.access(os.path.join(AtomDir,"POTCAR.gz"),os.R_OK):
            os.system("zcat %s >> %s"%(os.path.join(AtomDir,"POTCAR.gz"),os.path.join(CALCDIR,"POTCAR")))
            POTCARS.append((Atom,os.path.join(AtomDir,"POTCAR.gz")))
        else:
            raise Exception("POSCAR file \"POTCAR\" or \"POTCAR.Z\" or \"POTCAR.gz\" can not read")
    return POTCARS

def WriteINCAR(INCAR,VTValues,POTCARINFO,CALCDIR):
    fp=open(os.path.join(CALCDIR,"INCAR"),'w')
    fp.write("#This INCAR file was generated by VASPTEMPLATES Code provided by Zicong\n")
    fp.write("#https://gitee.com/xczics/vasptemplate#\n")
    fp.write("#The POTCAR used in this calculation are:\n")
    for Atom in POTCARINFO:
        fp.write("#\t%s\t%s\n"%(Atom[0],Atom[1]))
    fp.write("#Generated date:%s\n"%(datetime.datetime.now().strftime('%Y/%m/%d-%X')))
    fp.write("#---------------------------------------#\n\n")
    fp.writelines(INCAR)
    fp.close()